from ultralytics import YOLO
from .color import get_colors
from .image import resize_image
import numpy as np
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import os

def put_chinese_text(img, text, position, font_size=32, color=(255, 255, 255)):
    """在图片上绘制中文文本"""
    # 创建一个空白的PIL图片，大小和原图一样
    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    
    # 加载系统字体
    try:
        if os.name == 'nt':  # Windows
            font_path = "C:/Windows/Fonts/simhei.ttf"
        else:  # Linux/Mac
            font_path = "/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf"
        font = ImageFont.truetype(font_path, font_size)
    except:
        # 如果找不到系统字体，使用默认字体
        font = ImageFont.load_default()
    
    # 绘制文本
    draw.text(position, text, font=font, fill=color)
    
    # 转换回OpenCV格式
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

def process_frame(frame, seg_model, emotion_model):
    # 实例分割
    seg_results = seg_model(frame)
    detections = []
    person_count = 0
    colors = get_colors(10)

    # 获取原始图像的副本
    processed_frame = frame.copy()
    
    for r in seg_results:
        boxes = r.boxes
        if len(boxes) == 0:
            print("未检测到人脸")
            continue

        print(f"检测到 {len(boxes)} 个可能的人脸")  # 调试信息
        for j in range(len(boxes)):
            class_id = int(boxes.cls[j].item())
            conf = float(boxes.conf[j].item())
            print(f"检测框 {j}: class_id={class_id}, conf={conf}")  # 调试信息
            
            # 降低置信度阈值到0.3，增加检测灵敏度
            if class_id == 0 and conf > 0.3:
                # 获取边界框坐标
                x1, y1, x2, y2 = map(int, boxes.xyxy[j].cpu().numpy())
                
                # 确保坐标在图像范围内
                x1 = max(0, x1)
                y1 = max(0, y1)
                x2 = min(frame.shape[1], x2)
                y2 = min(frame.shape[0], y2)
                
                # 裁剪人脸区域
                person_image = frame[y1:y2, x1:x2]

                if person_image.size > 0:
                    print(f"检测到人脸区域: ({x1}, {y1}, {x2}, {y2})")
                    
                    # 调整图像大小用于情绪识别
                    person_image = resize_image(person_image, 640)
                    emotion_results = emotion_model(person_image)
                    print(f"情绪识别结果: {emotion_results}")  # 调试信息

                    # 获取情绪标签
                    if len(emotion_results[0].boxes) > 0:
                        emotion_box = emotion_results[0].boxes[0]
                        emotion_class_id = int(emotion_box.cls[0].item())
                        emotion_conf = float(emotion_box.conf[0].item())
                        emotion_names = emotion_model.names
                        emotion_label = emotion_names[emotion_class_id]
                        print(f"检测到的情绪: {emotion_label}, 置信度: {emotion_conf:.2f}")
                        
                        # 将英文情绪标签转换为中文
                        emotion_chinese = {
                            'Anger': '生气',
                            'Contempt': '蔑视',
                            'Disgust': '厌恶',
                            'Fear': '恐惧',
                            'Happy': '高兴',
                            'Neutral': '中性',
                            'Sad': '伤心',
                            'Surprise': '惊讶',
                            'Unknown': '未知'
                        }.get(emotion_label, '未知')
                        
                        # 绘制检测框
                        color = (0, 0, 255)  # 红色
                        cv2.rectangle(processed_frame, (x1, y1), (x2, y2), color, 2)
                        
                        # 添加情绪标签
                        label = f"{emotion_chinese} {emotion_conf:.2f}"
                        
                        # 绘制标签背景
                        font_size = 32
                        bg_height = font_size + 10
                        cv2.rectangle(processed_frame, 
                                    (x1, y1 - bg_height), 
                                    (x1 + 150, y1), 
                                    color, 
                                    -1)
                        
                        # 使用自定义函数绘制中文文本
                        processed_frame = put_chinese_text(
                            processed_frame,
                            label,
                            (x1 + 5, y1 - bg_height + 5),
                            font_size=font_size,
                            color=(255, 255, 255)  # 白色
                        )
                        
                    else:
                        emotion_label = 'Unknown'
                        emotion_conf = 0.0
                        print("未检测到情绪表情")

                    # 使用YOLO格式的检测框 [x1, y1, x2, y2]
                    detection = (
                        np.array([x1, y1, x2, y2]),  # 边界框坐标
                        conf,                         # 检测置信度
                        emotion_label                 # 情绪标签
                    )
                    detections.append(detection)
                    print(f"添加检测结果: 边界框={[x1, y1, x2, y2]}, 置信度={conf:.2f}, 情绪={emotion_label}")

    print(f"检测到 {len(detections)} 个人脸")
    return detections, processed_frame  # 返回检测结果和处理后的帧